#include <stdio.h>
#include <stdlib.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <cuComplex.h>
#include <time.h>

// 错误检查宏：检查CUDA API调用是否成功
#define CHECK_CUDA(call) { \
    cudaError_t err = call; \
    if (err != cudaSuccess) { \
        fprintf(stderr, "CUDA错误: %s (文件: %s, 行: %d)\n", cudaGetErrorString(err), __FILE__, __LINE__); \
        exit(EXIT_FAILURE); \
    } \
}

// 错误检查宏：检查CUBLAS API调用是否成功
#define CHECK_CUBLAS(call) { \
    cublasStatus_t status = call; \
    if (status != CUBLAS_STATUS_SUCCESS) { \
        fprintf(stderr, "CUBLAS错误: %d (文件: %s, 行: %d)\n", status, __FILE__, __LINE__); \
        exit(EXIT_FAILURE); \
    } \
}

int main() {
    // 定义矩阵维度和批处理大小
    const int M = 50;            // 矩阵行数
    const int N = 1604;          // 矩阵列数
    const int batchSize = 512 * 3; // 批处理数量：1536个矩阵
    const int numRuns = 1000;    // 运行次数

    // 定义复数1和0，用于CUBLAS计算
    const cuComplex cuComplexOne = make_cuComplex(1.0f, 0.0f);
    const cuComplex cuComplexZero = make_cuComplex(0.0f, 0.0f);

    cuComplex (*gpu_A)[512][1604][50];
    cuComplex (*gpu_B)[512][50][50];
    CHECK_CUDA(cudaMalloc((void **)&gpu_A, sizeof(cuComplex) * M * N * batchSize));
    CHECK_CUDA(cudaMalloc((void **)&gpu_B, sizeof(cuComplex) * M * N * batchSize));

    // 分配主机端内存，用于存储输入矩阵A
    cuComplex *h_A = (cuComplex*)malloc(sizeof(cuComplex) * M * N * batchSize);
    if (!h_A) {
        fprintf(stderr, "主机内存分配失败\n");
        exit(EXIT_FAILURE);
    }
    cuComplex *h_C = (cuComplex*)malloc(sizeof(cuComplex) * M * M * batchSize);
    if (!h_C) {
        fprintf(stderr, "主机内存分配失败\n");
        exit(EXIT_FAILURE);
    }

    // 生成随机复数矩阵
    for (int i = 0; i < batchSize * M * N; i++) {
        h_A[i].x = (float)rand() / RAND_MAX * 2.0f - 1.0f;  // 实部
        h_A[i].y = (float)rand() / RAND_MAX * 2.0f - 1.0f;  // 虚部
    }

    // 分配设备端内存
    cuComplex *d_A, *d_B, *d_C; // d_A: 输入矩阵, d_B: 中间结果, d_C: 逆矩阵

    CHECK_CUDA(cudaMalloc(&d_A, sizeof(cuComplex) * M * N * batchSize));
    CHECK_CUDA(cudaMalloc(&d_B, sizeof(cuComplex) * M * M * batchSize));
    CHECK_CUDA(cudaMalloc(&d_C, sizeof(cuComplex) * M * M * batchSize));

    // 将输入数据从主机拷贝到设备
    CHECK_CUDA(cudaMemcpy(d_A, h_A, sizeof(cuComplex) * M * N * batchSize, cudaMemcpyHostToDevice));

    // 创建CUBLAS句柄
    cublasHandle_t handle;
    CHECK_CUBLAS(cublasCreate(&handle));

    // 为批量操作分配指针数组
    cuComplex **d_A_array, **d_B_array, **d_C_array;
    CHECK_CUDA(cudaMalloc(&d_A_array, sizeof(cuComplex*) * batchSize));
    CHECK_CUDA(cudaMalloc(&d_B_array, sizeof(cuComplex*) * batchSize));
    CHECK_CUDA(cudaMalloc(&d_C_array, sizeof(cuComplex*) * batchSize));

    // 主机端指针数组，用于填充设备端指针
    cuComplex *h_A_ptrs[batchSize], *h_B_ptrs[batchSize], *h_C_ptrs[batchSize];
    for (int i = 0; i < batchSize; i++) {
        h_A_ptrs[i] = d_A + i * M * N;   // 每个输入矩阵的起始地址
        h_B_ptrs[i] = d_B + i * M * M;   // 每个中间矩阵的起始地址
        h_C_ptrs[i] = d_C + i * M * M;   // 每个逆矩阵的起始地址
    }

    // 将指针数组拷贝到设备
    CHECK_CUDA(cudaMemcpy(d_A_array, h_A_ptrs, sizeof(cuComplex*) * batchSize, cudaMemcpyHostToDevice));
    CHECK_CUDA(cudaMemcpy(d_B_array, h_B_ptrs, sizeof(cuComplex*) * batchSize, cudaMemcpyHostToDevice));
    CHECK_CUDA(cudaMemcpy(d_C_array, h_C_ptrs, sizeof(cuComplex*) * batchSize, cudaMemcpyHostToDevice));

    // 为LU分解分配辅助数组
    int *d_pivotArray, *d_infoArray;
    CHECK_CUDA(cudaMalloc(&d_pivotArray, sizeof(int) * M * batchSize));
    CHECK_CUDA(cudaMalloc(&d_infoArray, sizeof(int) * batchSize));

    // 时间统计变量
    float totalTime = 0.0f, maxTime = 0.0f, minTime = 1e9f;
    float gemmTime = 0.0f, invTime = 0.0f, copyTime = 0.0f;

    // 运行1000次进行性能测试
    for (int run = 0; run < numRuns; run++) {
        // 创建CUDA事件用于计时
        cudaEvent_t start, stop, gemmStart, gemmStop, invStart, invStop, copyStart, copyStop;
        CHECK_CUDA(cudaEventCreate(&start));
        CHECK_CUDA(cudaEventCreate(&stop));
        CHECK_CUDA(cudaEventCreate(&gemmStart));
        CHECK_CUDA(cudaEventCreate(&gemmStop));
        CHECK_CUDA(cudaEventCreate(&invStart));
        CHECK_CUDA(cudaEventCreate(&invStop));
        CHECK_CUDA(cudaEventCreate(&copyStart));
        CHECK_CUDA(cudaEventCreate(&copyStop));

        // 记录总时间起点
        CHECK_CUDA(cudaEventRecord(start, 0));

        // 步骤1：矩阵转置相乘 (A * A^T)
        CHECK_CUDA(cudaEventRecord(gemmStart, 0));

        CHECK_CUBLAS(cublasCgemmBatched(
            handle,                  // CUBLAS句柄
            CUBLAS_OP_N,             // A矩阵不转置（原样使用）
            CUBLAS_OP_T,             // B矩阵转置（A^T）
            M,                       // 结果矩阵C的行数（A的行数）
            M,                       // 结果矩阵C的列数（A^T的列数，即原A的行数）
            N,                       // 公共维度（A的列数，A^T的行数）
            &cuComplexOne,           // alpha值（复数1，用于标量乘法）
            (const cuComplex**)&gpu_A,  // 输入矩阵A的指针数组（设备内存）
            M,                       // A矩阵的前导维度（行数）
            (const cuComplex**)&gpu_A,  // 输入矩阵B的指针数组（此处复用A矩阵）
            M,                       // B矩阵的前导维度（转置后的行数）
            &cuComplexZero,           // beta值（复数0，用于标量乘法）
            &gpu_B,               // 输出矩阵C的指针数组（设备内存）
            M,                       // C矩阵的前导维度（行数）
            batchSize));             // 批量处理的数量（1536个矩阵）

        CHECK_CUDA(cudaEventRecord(gemmStop, 0));
        CHECK_CUDA(cudaEventSynchronize(gemmStop));
        float elapsedGemm;
        CHECK_CUDA(cudaEventElapsedTime(&elapsedGemm, gemmStart, gemmStop));
        gemmTime += elapsedGemm;

        // 步骤2：矩阵求逆
        CHECK_CUDA(cudaEventRecord(invStart, 0));
        // LU分解

        CHECK_CUBLAS(cublasCgetrfBatched(
            handle,          // CUBLAS上下文句柄
            M,               // 矩阵阶数（方阵的行/列数）
            d_B_array,       // 输入矩阵指针数组（设备内存，每个元素指向MxM矩阵）
            M,               // 矩阵的前导维度（存储的行数）
            d_pivotArray,    // 置换数组（设备内存，每个矩阵需要M个int存储行交换信息）
            d_infoArray,     // 状态信息数组（设备内存，每个矩阵返回状态码）
            batchSize));     // 批量处理的数量（1536个矩阵）

        // 计算逆矩阵
        CHECK_CUBLAS(cublasCgetriBatched(
            handle,                  // CUBLAS上下文句柄
            M,                      // 矩阵阶数（与LU分解保持一致）
            (const cuComplex**)d_B_array,  // 输入矩阵指针数组（LU分解后的矩阵）
            M,                      // 矩阵的前导维度（存储的行数）
            d_pivotArray,           // 置换数组（来自getrf的结果）
            d_C_array,              // 输出矩阵指针数组（存放逆矩阵结果）
            M,                      // 逆矩阵的前导维度（存储的行数）
            d_infoArray,            // 状态信息数组（检查逆运算是否成功）
            batchSize));            // 批量处理的数量（保持与分解一致）

        CHECK_CUDA(cudaEventRecord(invStop, 0));
        CHECK_CUDA(cudaEventSynchronize(invStop));
        float elapsedInv;
        CHECK_CUDA(cudaEventElapsedTime(&elapsedInv, invStart, invStop));
        invTime += elapsedInv;

        // 数据拷贝回主机（此处注释，实际使用时可启用）
        CHECK_CUDA(cudaEventRecord(copyStart, 0));
        CHECK_CUDA(cudaMemcpy(h_C, d_C, sizeof(cuComplex) * M * M * batchSize, cudaMemcpyDeviceToHost));
        CHECK_CUDA(cudaEventRecord(copyStop, 0));
        CHECK_CUDA(cudaEventSynchronize(copyStop));
        float elapsedCopy;
        CHECK_CUDA(cudaEventElapsedTime(&elapsedCopy, copyStart, copyStop));
        copyTime += elapsedCopy;

        // 记录总时间终点
        CHECK_CUDA(cudaEventRecord(stop, 0));
        CHECK_CUDA(cudaEventSynchronize(stop));
        float elapsedTotal;
        CHECK_CUDA(cudaEventElapsedTime(&elapsedTotal, start, stop));
        totalTime += elapsedTotal;
        if (elapsedTotal > maxTime) maxTime = elapsedTotal;
        if (elapsedTotal < minTime) minTime = elapsedTotal;

        // 实时打印每次运行的时间
        printf("运行 %d: 总时间 = %.3f ms, 矩阵乘法 = %.3f ms, 求逆 = %.3f ms, 拷贝 = %.3f ms\n",
               run, elapsedTotal, elapsedGemm, elapsedInv, elapsedCopy);

        // 销毁CUDA事件
        CHECK_CUDA(cudaEventDestroy(start));
        CHECK_CUDA(cudaEventDestroy(stop));
        CHECK_CUDA(cudaEventDestroy(gemmStart));
        CHECK_CUDA(cudaEventDestroy(gemmStop));
        CHECK_CUDA(cudaEventDestroy(invStart));
        CHECK_CUDA(cudaEventDestroy(invStop));
        CHECK_CUDA(cudaEventDestroy(copyStart));
        CHECK_CUDA(cudaEventDestroy(copyStop));
    }

    // 计算平均时间
    float avgTotalTime = totalTime / numRuns;
    float avgGemmTime = gemmTime / numRuns;
    float avgInvTime = invTime / numRuns;
    float avgCopyTime = copyTime / numRuns;

    // 打印统计信息
    printf("\n%d次运行的统计信息:\n", numRuns);
    printf("平均总时间: %.3f ms\n", avgTotalTime);
    printf("平均矩阵乘法时间: %.3f ms\n", avgGemmTime);
    printf("平均求逆时间: %.3f ms\n", avgInvTime);
    printf("平均拷贝时间: %.3f ms\n", avgCopyTime);
    printf("最大总时间: %.3f ms\n", maxTime);
    printf("最小总时间: %.3f ms\n", minTime);

    // 释放资源
    CHECK_CUBLAS(cublasDestroy(handle));
    CHECK_CUDA(cudaFree(d_A));
    CHECK_CUDA(cudaFree(d_B));
    CHECK_CUDA(cudaFree(d_C));
    CHECK_CUDA(cudaFree(d_A_array));
    CHECK_CUDA(cudaFree(d_B_array));
    CHECK_CUDA(cudaFree(d_C_array));
    CHECK_CUDA(cudaFree(d_pivotArray));
    CHECK_CUDA(cudaFree(d_infoArray));
    free(h_A);

    return 0;
}
